# Copyright 2022 The TensorStore Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""CMake implementation of Bazel build `rule` function.

https://bazel.build/rules/lib/globals#rule

Currently this just supports a very limited set of functionality.
"""

# pylint: disable=relative-beyond-top-level,missing-function-docstring,protected-access,invalid-name,g-importing-member,g-short-docstring-punctuation
from typing import Any, Callable

from .bazel_target import TargetId
from .common_providers import FilesProvider
from .ignored import IgnoredObject
from .invocation_context import InvocationContext
from .label import Label
from .label import RelativeLabel
from .provider import TargetInfo
from .rule_ctx import RuleCtx
from .rule_ctx import RuleFile
from .scope_build_file import ScopeBuildBzlFile
from .select import Configurable


class Attr:
  """Represents a defined rule attribute."""

  def __init__(self, handle):
    self._handle = handle


class AttrModule:
  """Defines rule attribute types."""

  @staticmethod
  def string(
      default: str = "",
      doc: str = "",
      mandatory: bool = False,
      values: list[str] | None = None,
  ):
    # https://bazel.build/rules/lib/attr#string
    del doc
    del values

    def handle(
        context: InvocationContext,
        name: str,
        value: str | None,
        outs: list[TargetId],
    ):
      if mandatory and value is None:
        raise ValueError(f"Attribute {name} not specified")
      if value is None:
        value = default

      del outs
      del context

      def impl(ctx: RuleCtx):
        setattr(ctx.attr, name, ctx._context.evaluate_configurable(value))

      return impl

    return Attr(handle)

  @staticmethod
  def string_list(
      mandatory: bool = False,
      allow_empty: bool = True,
      *,
      default: list[str] | None = None,
      doc: str = "",
  ):
    """https://bazel.build/rules/lib/attr#string_list"""
    del doc

    if default is None:
      default = []

    def handle(
        context: InvocationContext,
        name: str,
        value: list[str] | None,
        outs: list[TargetId],
    ):
      if mandatory and value is None:
        raise ValueError(f"Attribute {name} not specified")
      if value is None:
        value = default
      if not value and not allow_empty:
        raise ValueError(f"Attribute {name} is empty")

      del outs
      del context

      def impl(ctx: RuleCtx):
        setattr(
            ctx.attr,
            name,
            ctx._context.evaluate_configurable_list(value),
        )

      return impl

    return Attr(handle)

  @staticmethod
  def label(
      default: RelativeLabel | Callable[[RuleCtx], RelativeLabel] | None = None,
      doc: str = "",
      executable: bool = False,
      allow_files: Any = None,
      allow_single_file: Any = None,
      mandatory: bool = False,
      **kwargs,
  ):
    """https://bazel.build/rules/lib/attr#label"""
    del doc
    del executable
    del allow_files
    del allow_single_file
    del kwargs

    def handle(
        context: InvocationContext,
        name: str,
        value: RelativeLabel | Callable[[RuleCtx], RelativeLabel] | None,
        outs: list[TargetId],
    ):
      if mandatory and value is None:
        raise ValueError(f"Attribute {name} not specified")

      del outs
      del context

      def impl(ctx: RuleCtx):
        nonlocal value
        if value is None:
          if callable(default):
            value = default(ctx)
            assert value is not None  # default cannot return None.
          else:
            value = default

        if value is None:
          setattr(ctx.attr, name, None)
          return
        target_id = ctx._context.resolve_target_or_label(value)
        setattr(ctx.attr, name, ctx._context.get_target_info(target_id))

      return impl

    return Attr(handle)

  @staticmethod
  def label_list(
      allow_empty=True,
      *,
      default: (
          list[RelativeLabel | Callable[[RuleCtx], RelativeLabel]] | None
      ) = None,
      mandatory: bool = False,
      **kwargs,
  ):
    """https://bazel.build/rules/lib/attr#label_list"""
    del kwargs
    del allow_empty

    def handle(
        context: InvocationContext,
        name: str,
        value: list[RelativeLabel] | None,
        outs: list[TargetId],
    ):
      if mandatory and value is None:
        raise ValueError(f"Attribute {name} not specified")

      del outs
      del context

      def impl(ctx: RuleCtx):
        nonlocal value
        if not value:
          value = []
          if default:
            for x in default:
              if callable(x):
                v = x(ctx)
                assert v is not None  # default cannot return None.
                value.append(v)
              else:
                value.append(x)
        targets = ctx._context.resolve_target_or_label_list(value)

        setattr(
            ctx.attr,
            name,
            [ctx._context.get_target_info(target) for target in targets],
        )

      return impl

    return Attr(handle)

  @staticmethod
  def output(doc: str = "", mandatory: bool = False):
    """https://bazel.build/rules/lib/attr#output"""
    del doc

    def handle(
        context: InvocationContext,
        name: str,
        value: RelativeLabel | None,
        outs: list[TargetId],
    ):
      if mandatory and value is None:
        raise ValueError(f"Attribute {name} not specified")

      target = None
      if value is not None:
        target = context.resolve_target_or_label(value)
        outs.append(target)

      del outs
      del value
      del context

      def impl(ctx: RuleCtx):
        if target is None:
          setattr(ctx.attr, name, None)
          setattr(ctx.outputs, name, None)
          return
        path = str(ctx._context.get_generated_file_path(target))
        ctx._context.add_analyzed_target(
            target, TargetInfo(FilesProvider([path]))
        )
        setattr(ctx.outputs, name, RuleFile(path))

      return impl

    return Attr(handle)

  @staticmethod
  def output_list(
      allow_empty=True,
      *,
      mandatory: bool = False,
      **kwargs,
  ):
    """https://bazel.build/rules/lib/attr#output_list"""
    del kwargs
    del allow_empty

    def handle(
        context: InvocationContext,
        name: str,
        value: Configurable[list[RelativeLabel]] | None,
        outs: list[TargetId],
    ):
      if mandatory and value is None:
        raise ValueError(f"Attribute {name} not specified")
      if value is None:
        value = []

      del outs
      del context

      def impl(ctx: RuleCtx):
        targets = ctx._context.resolve_target_or_label_list(
            ctx._context.evaluate_configurable_list(value)
        )
        paths = []
        for x in targets:
          p = str(ctx._context.get_generated_file_path(x))
          ctx._context.add_analyzed_target(x, TargetInfo(FilesProvider([p])))
          paths.append(p)
        setattr(ctx.outputs, name, [RuleFile(p) for p in paths])

      return impl

    return Attr(handle)

  @staticmethod
  def bool_(default: bool = False, doc: str = "", mandatory: bool = False):
    # https://bazel.build/rules/lib/attr#bool
    del doc

    def handle(
        context: InvocationContext,
        name: str,
        value: bool | None,
        outs: list[TargetId],
    ):
      if mandatory and value is None:
        raise ValueError(f"Attribute {name} not specified")
      if value is None:
        value = default

      del outs
      del context

      def impl(ctx: RuleCtx):
        setattr(ctx.attr, name, ctx._context.evaluate_configurable(value))

      return impl

    return Attr(handle)


def _rule_impl(
    _context: InvocationContext,
    target: TargetId,
    attr_impls: list[Callable[[RuleCtx], None]],
    implementation: Callable[[RuleCtx], Any],
):
  ctx = RuleCtx(_context, Label(target, _context.workspace_root_for_label))
  for a in attr_impls:
    a(ctx)
  result = implementation(ctx)
  if result is None:
    result = []
  elif not isinstance(result, list):
    result = [result]
  _context.add_analyzed_target(target, TargetInfo(*result))


def rule(
    self: ScopeBuildBzlFile,
    implementation: Callable[[RuleCtx], Any],
    *,
    attrs: dict[str, Attr] | None = None,
    executable: bool = False,
    output_to_genfiles: bool = False,
    doc: str = "",
    **rule_kwargs,
):
  """https://bazel.build/rules/lib/globals#rule"""
  del executable
  del output_to_genfiles
  del doc

  def rule_func(
      name: str, visibility: list[RelativeLabel] | None = None, **kwargs
  ):
    # snaptshot the invocation context.
    context = self._context.snapshot()
    target = context.resolve_target(name)

    outs: list[TargetId] = []

    if attrs is not None:
      attr_impls = [
          attr_obj._handle(
              context, attr_name, kwargs.pop(attr_name, None), outs
          )
          for attr_name, attr_obj in attrs.items()
      ]
    else:
      attr_impls = []

    context.add_rule(
        target,
        lambda: _rule_impl(context, target, attr_impls, implementation),
        outs=outs,
        visibility=visibility,
    )

  return rule_func


setattr(AttrModule, "bool", AttrModule.bool_)


setattr(ScopeBuildBzlFile, "bazel_attr", AttrModule)
setattr(ScopeBuildBzlFile, "bazel_rule", rule)
setattr(ScopeBuildBzlFile, "bazel_CcInfo", IgnoredObject())
setattr(ScopeBuildBzlFile, "bazel_ProtoInfo", IgnoredObject())
setattr(ScopeBuildBzlFile, "bazel_DefaultInfo", IgnoredObject())
